#include "kernel_operator.h"

#include "my_functions.h"
#include "op_common.h"
#include <cstdint>
#include <type_traits>

using namespace AscendC;

template <typename _DT_X, typename _DT_Y> class Kernel
{
  public:
    TPipe pipe;
    CommonTiling tilingData;
    int64_t finalDim;
    bool largeFinal;
    int64_t alignSize;
    DefGlobalTensor(X);
    DefInQue(X0);
    DefInQue(X1);
    DefOutTensor(Y);

    DefBufVECIN(CASTED_X0);
    DefBufVECIN(CASTED_X1);
    DefBufVECIN(TEMP_Y);

  public:
    __aicore__ inline Kernel()
    {
    }

    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, CommonTiling &tiling, int64_t final_dim, bool large_final)
    {
        tilingData = tiling;
        finalDim = final_dim;
        largeFinal = large_final;
        alignSize = 32 / sizeof(_DT_X);
        // 切分global
        GSetBuffer(X, x, tilingData.startIdx, tilingData.blockLength);
        GSetBuffer(Y, y, tilingData.startIdx, tilingData.blockLength / 2);
        // 初始化队列
        auto queLength = ALIGN_TO(tilingData.bufferSize * sizeof(_DT_X), 32);
        InitQueue(X0, queLength);
        InitQueue(X1, queLength);
        InitQueue(Y, queLength);
        if constexpr (std::is_same_v<_DT_X, half>)
        {
            queLength *= 2;
            InitTBufBuffer(CASTED_X0, queLength);
            InitTBufBuffer(CASTED_X1, queLength);
            InitTBufBuffer(TEMP_Y, queLength);
        }
    }

    __aicore__ inline void ProcessForLargeFinal()
    {
        uint32_t halfFinalDim = finalDim / 2;
        uint32_t loopCount = tilingData.blockLength / finalDim;
        uint32_t loopCountFinalDim = CEIL_DIV(halfFinalDim, tilingData.bufferSize);
        uint32_t finalJ = loopCountFinalDim - 1;
        for (auto i = 0; i < loopCount; ++i)
        {
            for (auto j = 0; j < finalJ; ++j)
            {
                auto index0 = i * finalDim + j * tilingData.bufferSize;
                auto index1 = index0 + halfFinalDim;
                auto indexy = i * halfFinalDim + j * tilingData.bufferSize;
                CopyInForLargeFinal(index0, index1, tilingData.bufferSize);
                Compute(tilingData.bufferSize);
                CopyOutForLargeFinal(indexy, tilingData.bufferSize);
            }
            auto copyCount = ALIGN_TO(tilingData.finalCalcCount, alignSize);
            auto index0 = i * finalDim + finalJ * tilingData.bufferSize;
            auto index1 = index0 + halfFinalDim;
            auto indexy = i * halfFinalDim + finalJ * tilingData.bufferSize;
            CopyInForLargeFinal(index0, index1, copyCount);
            Compute(tilingData.finalCalcCount);
            CopyOutForLargeFinalEnd(indexy, tilingData.finalCalcCount);
        }
    }

    __aicore__ inline void Process()
    {
        uint32_t halfFinalDim = finalDim / 2;
        auto halfFinalCount = tilingData.blockLength / 2 / halfFinalDim;
        auto tileFinalCount = tilingData.bufferSize / halfFinalDim;
        auto loopCount = CEIL_DIV(halfFinalCount, tileFinalCount);
        auto finnal_progress = loopCount - 1;
        auto copyCount = ALIGN_TO(tilingData.bufferSize, alignSize);
        for (uint32_t i = 0; i < finnal_progress; ++i)
        {
            CopyIn(i, tilingData.bufferSize);
            Compute(tilingData.bufferSize);
            CopyOut(i, copyCount);
        }
        CopyIn(finnal_progress, tilingData.finalCalcCount);
        Compute(tilingData.finalCalcCount);
        CopyOutFinal(finnal_progress, tilingData.finalCalcCount);
    }

    __aicore__ inline void CopyIn(uint32_t i, uint32_t calcCount)
    {
        uint32_t halfFinalDim = finalDim / 2;
        auto copyCount = ALIGN_TO(halfFinalDim, alignSize);
        QueAlloc(X0, TypeOf(X));
        QueAlloc(X1, TypeOf(X));
        for (auto j = 0; j < calcCount / halfFinalDim; ++j)
        {
            auto localIndex = j * halfFinalDim;
            auto index0 = i * tilingData.bufferSize * 2 + j * finalDim;
            auto index1 = index0 + halfFinalDim;
            DataCopy(LTensorName(X0)[localIndex], GTensorName(X)[index0], copyCount);
            DataCopy(LTensorName(X1)[localIndex], GTensorName(X)[index1], copyCount);
        }
        EnQue(X0);
        EnQue(X1);
    }

    __aicore__ inline void CopyOut(uint32_t i, uint32_t calcCount)
    {
        auto index = i * tilingData.bufferSize;
        DeQueLocal2Global(Y, index, calcCount);
    }

    __aicore__ inline void CopyOutFinal(uint32_t i, uint32_t calcCount)
    {
        auto index = i * tilingData.bufferSize;
        DeQueSimple(Y);
        if (tilingData.bufferSize < alignSize)
        {
            for (auto j = 0; j < calcCount; ++j)
            {
                GTensorName(Y).SetValue(index + j, LTensorName(Y).GetValue(j));
            }
        }
        else
        {
            auto copyCount = ALIGN_TO(calcCount, alignSize);
            auto index2 = calcCount / alignSize * alignSize;
            // 自动忽略后面不足alignSize的数据
            DataCopy(GTensorName(Y)[index], LTensorName(Y), calcCount);
            // 最后32B
            if (index2 != calcCount)
            {
                DataCopy(GTensorName(Y)[index + index2], LTensorName(Y)[index2], alignSize);
            }
        }
        QueFree(Y);
    }

    __aicore__ inline void CopyInForLargeFinal(uint32_t i, uint32_t j, uint32_t calcCount)
    {
        QueAlloc(X0, TypeOf(X));
        DataCopy(LTensorName(X0), GTensorName(X)[i], calcCount);
        EnQue(X0);
        QueAlloc(X1, TypeOf(X));
        DataCopy(LTensorName(X1), GTensorName(X)[j], calcCount);
        EnQue(X1);
    }

    __aicore__ inline void Compute(uint32_t calcCount)
    {
        if constexpr (std::is_same_v<_DT_X, half>)
        {
            ComputeHalf(calcCount);
        }
        else
        {
            DeQue(X0, TypeOf(X));
            DeQue(X1, TypeOf(X));
            QueAllocSimple(Y);
            my_geglu(LTensorName(Y), LTensorName(X0), LTensorName(X1), calcCount);
            EnQue(Y);
            QueFree(X0);
            QueFree(X1);
        }
    }

    __aicore__ inline void ComputeHalf(uint32_t calcCount)
    {
        DeQue(X0, TypeOf(X));
        DeQue(X1, TypeOf(X));
        QueAllocSimple(Y);
        TBufGet(CASTED_X0, float);
        TBufGet(CASTED_X1, float);
        TBufGet(TEMP_Y, float);
        Cast(TBufTensorName(CASTED_X0), LTensorName(X0), RoundMode::CAST_NONE, calcCount);
        Cast(TBufTensorName(CASTED_X1), LTensorName(X1), RoundMode::CAST_NONE, calcCount);
        my_geglu(TBufTensorName(TEMP_Y), TBufTensorName(CASTED_X0), TBufTensorName(CASTED_X1), calcCount);
        Cast(LTensorName(Y), TBufTensorName(TEMP_Y), RoundMode::CAST_ROUND, calcCount);
        EnQue(Y);
        QueFree(X0);
        QueFree(X1);
    }

    __aicore__ inline void CopyOutForLargeFinal(uint32_t i, uint32_t copyCount)
    {
        DeQueLocal2Global(Y, i, copyCount);
    }

    __aicore__ inline void CopyOutForLargeFinalEnd(uint32_t i, uint32_t calcCount)
    {
        DeQueSimple(Y);
        DataCopy(GTensorName(Y)[i], LTensorName(Y), calcCount);
        auto index2 = calcCount / alignSize * alignSize;
        if (index2 != calcCount)
        {
            DataCopy(GTensorName(Y)[i + index2], LTensorName(Y)[index2], alignSize);
        }
        QueFree(Y);
    }
};

extern "C" __global__ __aicore__ void geglu(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    CRATE_COMMON_TILING_SIMPLE;
    Kernel<DTYPE_X, DTYPE_Y> op;
    op.Init(x, y, commonTiling, tiling_data.finalDim, tiling_data.largeFinal);
    if (tiling_data.largeFinal)
    {
        op.ProcessForLargeFinal();
    }
    else
    {
        op.Process();
    }
}